//*********************************************
//  This is Geant4 Template
//                                  author:Qian
//

#include "SimEvent.h"
#include "Verbose.hh"

#include "G4SystemOfUnits.hh"
#include "G4Box.hh"
#include "G4VPhysicalVolume.hh"
#include "G4LogicalVolumeStore.hh"
#include "G4PhysicalVolumeStore.hh"

#include "MyDetectorConstruction.hh"
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wshadow"
#include "TObject.h"
#include "TFile.h"
#include "TTree.h"
#include "TSystem.h"
#include "TNtuple.h"
#include "TH1.h"
#include "TH2.h"
#include "TVector3.h"

#pragma GCC diagnostic pop

#include "MyGenContainer.hh"
#include "MyRootBasedAnalysis.hh"

//....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......

MyRootBasedAnalysis::MyRootBasedAnalysis()
{
    SetFileName("output.root");
    m_EdepInSD = 0.;
    m_PhoElec = false;
    fTotalEventNumber = 0;
    fGenValidEvents = 0;
}

//....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......

MyRootBasedAnalysis::~MyRootBasedAnalysis()
{
}

//....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......

void MyRootBasedAnalysis::BeginOfRunAction()
{
    if (!active)
        return;

    if (verbose)
        G4cout << "====>MyRootBasedAnalysis::BeginOfRunAction()" << G4endl;

    fRootFp = new TFile(fFileName, "recreate");
    if (!fRootFp)
    {
        G4cout << "\n====>MyRootBasedAnalysis::BeginOfRunAction(): "
               << "cannot open " << fFileName << G4endl;
        return;
    }

    m_GenPtr = MyGenContainer::GetInstance().GetMyGen();

    fEvent = new SimEvent();
    fConfigTree = new TTree("MCConfig", "Tree of ParticleGun Config");
    fConfigTree->Branch("GitVersion", &fGitVersion);
    fConfigTree->Branch("PGType", &fPGType);
    fConfigTree->Branch("GenTime", &fGenTime, "GenTime/D");
    fConfigTree->Branch("GenEvents", &fGenEvents, "GenEvents/I");
    fConfigTree->Branch("GenValidEvents", &fGenValidEvents, "GenValidEvents/I");
    fConfigTree->Branch("TotalEventNumber", &fTotalEventNumber, "TotalEventNumber/I");
    fConfigTree->Branch("SDSize", &fSDSize, "SDSize[3]/D");
    fConfigTree->Branch("NumOfSD", &fNumOfSD, "NumOfSD/i");
    fConfigTree->Branch("SDPosition", &fSDPosition, "SDPosition[NumOfSD][3]/D");

    fTree = new TTree("Sim", "Tree of data events");
    TBranch* br_SimEvent = fTree->Branch("SimEvent", "SimEvent", fEvent, 32000, 0);
    br_SimEvent->SetAutoDelete(true);

    //------- add your codes down here
    //
    return;
}

//....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......

void MyRootBasedAnalysis::EndOfRunAction()
{
    if (!active)
        return;

    if (verbose)
        G4cout << "====>MyRootBasedAnalysis::EndOfRunAction()" << G4endl;

    if (!fRootFp)
    {
        G4cout << "\n====>MyRootBasedAnalysis::EndOfRunAction(): "
               << "cannot open " << fFileName << G4endl;
        return;
    }

    //------- add your codes down here
    //

    YAML::Node yamlConfig = MyGenContainer::GetInstance().GetConfig();
    //fGitVersion = TString(VERSION_GIT_HEAD_VERSION);
    fPGType = yamlConfig["Type"].as<std::string>("SimpleParticleGun");
    fGenEvents = yamlConfig["EventNumber"].as<int>(100);
    this->GetSDInfo();


    fConfigTree->Fill();
    fRootFp->cd();
    fConfigTree->AutoSave();

    G4cout << "\n====>In total " << fTree->GetEntries() << " Events have been stored." << G4endl;
    G4cout << "Totally " << fTotalEventNumber << " events have been generated." << G4endl;
    G4cout << "Generator level cut efficiency: " << fGenValidEvents / double(fTotalEventNumber) << G4endl;
    fRootFp->cd();
    fTree->AutoSave();
    //fRootFp->Write();
    fRootFp->Close();
    return;
}

//....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......

void MyRootBasedAnalysis::BeginOfEventAction(const G4Event *)
{
    if (!active)
        return;

    if (verbose > 1)
        G4cout << "====>MyRootBasedAnalysis::BeginOfEventAction()" << G4endl;

    m_EdepInSD = 0.;
    m_PhoElec = false;
    fEvent->MyClear();

    //------- add your codes down here
    //
    return;
}

//....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......

void MyRootBasedAnalysis::EndOfEventAction(const G4Event *)
{
    if (!active)
        return;

    if (verbose > 1)
        G4cout << "====>MyRootBasedAnalysis::EndOfEventAction()" << G4endl;

    //------- add your codes down here
    //
    ++fTotalEventNumber;
    if(m_ValidGen && m_EdepInSD <= 0.)
    {
        m_GenPtr->SetValid(false);
        fEvent->MyClear();
        return;
    }
    
    if(m_ValidGen && m_GenPtr->IfNeedPolPhot() && !m_PhoElec)
    {
        m_GenPtr->SetValid(false);
        fEvent->MyClear();
        return;
    }

    ++fGenValidEvents;
    if(fGenValidEvents % 1000 == 0)
        std::cout << "ValidEvent: " << fGenValidEvents << ", TotalEvent: " << fTotalEventNumber << std::endl;
    m_GenPtr->SetValid(true);
    fEvent->SetTruthParticle(m_GenPtr->GetParticle().data());
    fEvent->SetTruthEnergy(m_GenPtr->GetParticleEnergy() / MeV);
    G4ThreeVector TruthPos = m_GenPtr->GetParticlePosition() / mm;
    G4ThreeVector TruthPol = m_GenPtr->GetParticlePolarization();
    G4ThreeVector TruthMomDir = m_GenPtr->GetParticleMomentumDirection();
    fEvent->SetTruthPosition(TVector3(TruthPos[0], TruthPos[1], TruthPos[2]));
    fEvent->SetTruthPolarization(TVector3(TruthPol[0], TruthPol[1], TruthPol[2]));
    fEvent->SetTruthMomentumDirection(TVector3(TruthMomDir[0], TruthMomDir[1], TruthMomDir[2]));

    fTree->Fill();
    fEvent->MyClear();

    return;
}

//....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......
G4ClassificationOfNewTrack MyRootBasedAnalysis::ClassifyNewTrack(const G4Track *)
{
    if (!active)
        return fUrgent;

    //------- add your codes down here
    //
    return fUrgent;
}

//....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......
void MyRootBasedAnalysis::PreTrackingAction(const G4Track *aTrack)
{
    if (!active)
        return;

    //------- add your codes down here
    //
    G4ParticleDefinition *particle = aTrack->GetDefinition();
    G4int trkID = aTrack->GetTrackID();
    SimTrack* aSimTrack = fEvent->GetTrack(trkID);
    aSimTrack->SetPDGID(particle->GetPDGEncoding());
    aSimTrack->SetTrackID(trkID);
    aSimTrack->SetParentID(aTrack->GetParentID());
    aSimTrack->SetInitMass(particle->GetAtomicMass() / MeV);
    aSimTrack->SetInitEk(aTrack->GetKineticEnergy() / MeV);

    G4ThreeVector InitMom = aTrack->GetMomentum();
    G4ThreeVector InitPos = aTrack->GetPosition();
    aSimTrack->SetInitMom(TVector3(InitMom[0] / MeV, InitMom[1] / MeV, InitMom[2] / MeV));
    aSimTrack->SetInitPos(TVector3(InitPos[0] / mm, InitPos[1] / mm, InitPos[2] / mm));
    aSimTrack->SetInitT(aTrack->GetGlobalTime() / s);
    //std::cout << aTrack->GetCreatorModelName() << std::endl;
    aSimTrack->SetCreatorProcess(aTrack->GetCreatorModelName().data());

    //delete aSimTrack;
    return;
}

//....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......
void MyRootBasedAnalysis::PostTrackingAction(const G4Track *aTrack)
{
    if (!active)
        return;

    //------- add your codes down here
    //
    SimTrack* aSimTrack = fEvent->GetTrack(aTrack->GetTrackID());
    
    G4ThreeVector ExitMom = aTrack->GetMomentum();
    G4ThreeVector ExitPos = aTrack->GetPosition();
    aSimTrack->SetExitMom(TVector3(ExitMom[0] / MeV, ExitMom[1] / MeV, ExitMom[2] / MeV));
    aSimTrack->SetExitPos(TVector3(ExitPos[0] / mm, ExitPos[1] / mm, ExitPos[2] / mm));
    aSimTrack->SetExitT(aTrack->GetGlobalTime() / s);

    aSimTrack->SetTrackLength(aTrack->GetTrackLength() / mm);
    aSimTrack->SetEdep((aSimTrack->GetInitEk() - aTrack->GetKineticEnergy()) / MeV);

    //delete aSimTrack;
    return;
}

//....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......
void MyRootBasedAnalysis::SteppingAction(const G4Step * aStep)
{
    if (!active)
        return;

    //------- add your codes down here
    //
    G4Track *aTrack = aStep->GetTrack();
    const G4ParticleDefinition *particle = aTrack->GetParticleDefinition();
    G4int particleID = particle->GetPDGEncoding();

    G4StepPoint *preStepPoint = aStep->GetPreStepPoint();
    G4StepPoint *postStepPoint = aStep->GetPostStepPoint();

    G4ThreeVector prePos = preStepPoint->GetPosition();
    G4ThreeVector preMomtum = preStepPoint->GetMomentum();

    G4ThreeVector postPos = postStepPoint->GetPosition();
    G4ThreeVector postMomtum = postStepPoint->GetMomentum();
    G4String postVolumeName = preStepPoint->GetPhysicalVolume()->GetName();

    G4String proName = postStepPoint->GetProcessDefinedStep()->GetProcessName();
    //G4int proSubType = postStepPoint->GetProcessDefinedStep()->GetProcessSubType();
    //G4ProcessType proType = postStepPoint->GetProcessDefinedStep()->GetProcessType();
    //G4int proSubType = postStepPoint->GetProcessDefinedStep()->GetProcessSubType();
    //G4String proTypeName = postStepPoint->GetProcessDefinedStep()->GetProcessTypeName(postStepPoint->GetProcessDefinedStep()->GetProcessType());

    //std::cout << proName << ", " << proTypeName << ", " << postStepPoint->GetProcessDefinedStep()->GetProcessSubType() << std::endl;

    G4double StepEdep = aStep->GetTotalEnergyDeposit();

    //if (particleID != 22 && postVolumeName == G4String("GasSDPhys") && StepEdep > 0.) fIfValidEvent = true; //标记有效事例
    if (proName == "pol-phot") m_PhoElec = true; //标记光电效应
    if(postVolumeName == G4String("GasSDPhys"))
    {
        m_EdepInSD += StepEdep;
    }

    //Save a Deposit
    SimDeposit* aSimDeposit = new SimDeposit();

    aSimDeposit->SetPDGID(particle->GetPDGEncoding());
    aSimDeposit->SetTrackID(aTrack->GetTrackID());
    aSimDeposit->SetParentID(aTrack->GetParentID());
    aSimDeposit->SetCharge(particle->GetPDGCharge());

    aSimDeposit->SetPreMomentum(TVector3(preMomtum[0] / MeV, preMomtum[1] / MeV, preMomtum[2] / MeV));
    aSimDeposit->SetPrePosition(TVector3(prePos[0] / mm, prePos[1] / mm, prePos[2] / mm));
    aSimDeposit->SetPreT(preStepPoint->GetGlobalTime() / s);
    aSimDeposit->SetVolumeName(postVolumeName.data());


    aSimDeposit->SetPostMomentum(TVector3(postMomtum[0] / MeV, postMomtum[1] / MeV, postMomtum[2] / MeV));
    aSimDeposit->SetPostPosition(TVector3(postPos[0] / mm, postPos[1] / mm, postPos[2] / mm));
    aSimDeposit->SetPostT(postStepPoint->GetGlobalTime() / s);
    aSimDeposit->SetProcessName(TString(proName.data()));
    //aSimDeposit->SetProcessSubType(proSubType);

    aSimDeposit->SetEdep(StepEdep / MeV);
    aSimDeposit->SetStepLength(aStep->GetStepLength() / mm);
    aSimDeposit->SetFirstDeposit(aStep->IsFirstStepInVolume());

    fEvent->AddDeposit(aSimDeposit->GetTrackID(), aSimDeposit);

    //delete aSimDeposit;

    if (particleID == 2212 && aTrack->GetKineticEnergy() < 10 * keV)
    {
        //std::cout << "Proton energy less than 8 keV, kill" << std::endl;
        aTrack->SetTrackStatus(fStopAndKill);
    }

    return;
}

void MyRootBasedAnalysis::GetSDInfo(G4String NameOfSD)
{
    G4LogicalVolumeStore* LogVolStore = G4LogicalVolumeStore::GetInstance();
    G4LogicalVolume* SDLogVol = LogVolStore->GetVolume(NameOfSD);

    //Get SD Size
    G4Box* SDSolVol = dynamic_cast<G4Box*>(SDLogVol->GetSolid());
    fSDSize[0] = SDSolVol->GetXHalfLength() * 2. / mm;
    fSDSize[1] = SDSolVol->GetYHalfLength() * 2. / mm;
    fSDSize[2] = SDSolVol->GetZHalfLength() * 2. / mm;
    std::cout << "SD size: X: " << fSDSize[0] << "\tY: " << fSDSize[1] << "\tZ: " << fSDSize[2] << std::endl;

    fNumOfSD = 0;
    G4PhysicalVolumeStore* PhyVolStore = G4PhysicalVolumeStore::GetInstance();
    G4ThreeVector GasSDPos;
    for (auto iPhyVol = PhyVolStore->begin(); iPhyVol != PhyVolStore->end(); ++iPhyVol)
    {
        G4String Name = (*iPhyVol)->GetName();
        if(Name.find(G4String("GasSDPhys")) != G4String::npos)
        {
            GasSDPos = (*iPhyVol)->GetTranslation() / mm;
            if(verbose) std::cout << "GasSDPos: X: " << fSDPosition[fNumOfSD][0] << "\tY: " << fSDPosition[fNumOfSD][1] << "\tZ: " << fSDPosition[fNumOfSD][2] << std::endl;
        }
    }
    for (auto iPhyVol = PhyVolStore->begin(); iPhyVol != PhyVolStore->end(); ++iPhyVol)
    {
        G4String Name = (*iPhyVol)->GetName();
        if(Name.find(G4String("DetectorPhys")) != G4String::npos)
        {
            G4ThreeVector Pos = (*iPhyVol)->GetTranslation() / mm;
            fSDPosition[fNumOfSD][0] = Pos[0] + GasSDPos[0];
            fSDPosition[fNumOfSD][1] = Pos[1] + GasSDPos[1];
            fSDPosition[fNumOfSD][2] = Pos[2] + GasSDPos[2];
            if(verbose) std::cout << "Catch SD, Pos: X: " << fSDPosition[fNumOfSD][0] << "\tY: " << fSDPosition[fNumOfSD][1] << "\tZ: " << fSDPosition[fNumOfSD][2] << std::endl;
            ++fNumOfSD;
        }
    }
    if(verbose) std::cout << "Catch totally " << fNumOfSD << " SD." << std::endl;
}